import rpy2.robjects as robjects
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
from utils2 import rotate1, rotate, move_y


def airfoil_show(title, ordinates, ordinates_low, ordinates_up, range_limit=True):
    plt.title(title)
    plt.scatter(ordinates[:, 0], ordinates[:, 1], color='b', s=2)
    plt.scatter(ordinates_low[:, 0], ordinates_low[:, 1], color='r', s=2)
    plt.scatter(ordinates_up[:, 0], ordinates_up[:, 1], color='r', s=2)
    plt.grid()  # 设置网格

    if range_limit:
        plt.xlim((-1.05, 1.05))  # 设置横坐标范围
        plt.ylim((-0.3, 0.3))  # 设置纵坐标范围
        plt.xticks(np.arange(-1.0, 1.05, 0.1))
    plt.gca().set_aspect(1)
    plt.show()


def airfoil_show2(title, ordinates, re_ordinates):
    plt.title(title)
    plt.scatter(ordinates[:, 0], ordinates[:, 1], color='b', label="origin")
    plt.scatter(re_ordinates[:, 0], re_ordinates[:, 1], color='r', label="recover")
    plt.grid()  # 设置网格

    plt.xlim((-0.05, 1.05))  # 设置横坐标范围
    plt.ylim((-0.3, 0.3))  # 设置纵坐标范围
    plt.xticks(np.arange(0, 1.05, 0.1))

    plt.legend()

    plt.gca().set_aspect(1)
    plt.show()


def x_to_arc(ordinates_up, ordinates_low, isShow=False):
    '''
        将翼型的横坐标用某点到原点的弧长代替
    :param x:
    :return:
    '''

    arc_ordinates_up = np.ones_like(ordinates_up)
    arc_ordinates_low = np.ones_like(ordinates_low)

    # 计算上机翼弧长
    arc_len = 0
    prv = ordinates_up[0]
    for i, ord in enumerate(ordinates_up[1::]):
        arc_len += np.sqrt((ord[0] - prv[0]) ** 2 + (ord[1] - prv[1]) ** 2)
        prv = ord
        arc_ordinates_up[1 + i, 0] = arc_len

    arc_ordinates_up[0, 0] = 0
    arc_ordinates_up[:, 1] = ordinates_up[:, 1]
    # print(arc_ordinates_up[0:5])

    # 计算下机翼弧长
    arc_len = 0
    prv = ordinates_low[0]
    for i, ord in enumerate(ordinates_low[1::]):
        arc_len += np.sqrt((ord[0] - prv[0]) ** 2 + (ord[1] - prv[1]) ** 2)
        prv = ord
        arc_ordinates_low[1 + i, 0] = arc_len

    arc_ordinates_low[0, 0] = 0
    arc_ordinates_low[:, 1] = ordinates_low[:, 1]
    # print(arc_ordinates_low[0:5])

    # 弧长开方
    arc_ordinates_up[:, 0] = np.sqrt(arc_ordinates_up[:, 0])
    arc_ordinates_low[:, 0] = np.sqrt(arc_ordinates_low[:, 0])

    # 下机翼取负
    arc_ordinates_low[:, 0] = arc_ordinates_low[:, 0] * -1

    # concatenate
    arc_ordinates = np.concatenate([arc_ordinates_low[::-1][0:-1], arc_ordinates_up], axis=0)
    # print("arc_ordinates:")
    # print(arc_ordinates)

    # show
    if isShow:
        plt.scatter(arc_ordinates[:, 0], arc_ordinates[:, 1], c='b')
        # plt.scatter(ordinates_low[:, 0], ordinates_low[:, 1], c='r')
        # plt.scatter(ordinates_up[:, 0], ordinates_up[:, 1], c='r')
        plt.show()

    return arc_ordinates_up, arc_ordinates_low, arc_ordinates


def produce_x_v1(N=65, isUp=True):
    '''
        产生余弦分布的x序列(上表面65个点；下表面不包括前缘顶点，为64个点)
    :return:
    '''
    phi = np.pi / 2 / (N - 1)

    x = 1 - np.cos(np.arange(0, N, 1) * phi)
    x[-1] = 1
    reverse_x = x[::-1]

    if isUp:
        # plot test
        # plt.title("up:reverse")
        # plt.scatter(np.arange(0, reverse_x.shape[0]), reverse_x)
        # plt.show()
        return reverse_x
    else:
        # plot test
        # plt.title("low:x")
        # plt.scatter(np.arange(0, x[:-1].shape[0]), x[:-1])
        # plt.show()
        return x[:-1]


def produce_x_v2(N=151, isUp=True):
    '''
        产生余弦分布的x序列(上表面151(N)个点；下表面不包括前缘顶点，为150(N-1)个点。另外，上下表面都必须包括后缘点)
    :return:
    '''

    tN = 2 * N - 1
    x = np.power((1 - np.cos((np.arange(0, tN, 1) - (N - 1)) * np.pi / (N - 1))), 2) / 4
    # plt.scatter(np.arange(0, x.shape[0]), x)
    # plt.show()

    if isUp:
        # plot test
        # plt.title("up")
        # plt.scatter(np.arange(0, x[0:N].shape[0]), x[0:N])
        # plt.show()
        return x[0:N]
    else:
        # plot test
        # plt.title("low")
        # plt.scatter(np.arange(0, x[N:2 * N].shape[0]), x[N:2 * N])
        # plt.show()
        return x[N - 1:2 * N - 1]


def ssp(x_train, y_train, isUp=True):
    x_smooth = produce_x_v2(isUp=isUp)
    r_y = robjects.FloatVector(y_train)
    r_x = robjects.FloatVector(x_train)
    # w = robjects.FloatVector(3 * np.ones(x_train.shape[0]))  # 所有点的权重相同
    # w = robjects.FloatVector(1.0 * np.power(np.e, -0.2 * x_train))  # 指数下降函——急
    # w = robjects.FloatVector(-1.0 * np.power(np.e, 0.08 * x_train) + np.e)  # 指数下降函数——缓
    w = robjects.FloatVector(1.0 * np.cos(0.1 * x_train))  # 余弦分布
    # w = robjects.FloatVector(1.0 - 0.25 * np.power(x_train, 2))  # 二次函数分布

    r_smooth_spline = robjects.r['smooth.spline']  # extract R function# run smoothing function
    # spline1 = r_smooth_spline(x=r_x, y=r_y, w=w, spar=0.8)
    spline1 = r_smooth_spline(x=r_x, y=r_y, spar=0.8)
    ySpline = np.array(robjects.r['predict'](spline1, robjects.FloatVector(x_smooth)).rx2('y'))
    # print(x_smooth.shape, ySpline.shape, x_smooth[0], x_smooth[-1], ySpline[0], ySpline[-1])

    # plot
    # plt.plot(x_smooth, ySpline)
    # plt.scatter(x_smooth, ySpline)

    if isUp:
        print(x_smooth)

    return np.concatenate([np.expand_dims(x_smooth, axis=1),
                           np.expand_dims(ySpline, axis=1)], axis=1)


def sspv2(x_train, y_train):
    x_smooth = x_train
    r_y = robjects.FloatVector(y_train)
    r_x = robjects.FloatVector(x_train)
    w = robjects.FloatVector(3 * np.ones(x_train.shape[0]))  # 所有点的权重相同

    r_smooth_spline = robjects.r['smooth.spline']  # extract R function# run smoothing function
    spline1 = r_smooth_spline(x=r_x, y=r_y, w=w, spar=0.8)
    ySpline = np.array(robjects.r['predict'](spline1, robjects.FloatVector(x_smooth)).rx2('y'))

    # plot
    # plt.plot(x_smooth, ySpline)
    # plt.scatter(x_smooth, ySpline)

    return np.concatenate([np.expand_dims(x_smooth, axis=1),
                           np.expand_dims(ySpline, axis=1)], axis=1)


def file_to_ssp(dir, file):
    file_path = os.path.join(dir, file)
    ordinates = pd.read_csv(file_path).to_numpy()

    # smooth.spline
    # 前缘顶点
    key = np.argmin(ordinates[:, 0])

    # 上表面坐标点
    ordinates_up = ordinates[0:key + 1]
    ordinates_up = ordinates_up[::-1]

    # 下表面坐标点
    ordinates_low = ordinates[key::]

    # 原翼型形状展示
    # airfoil_show("origin", ordinates, ordinates_low, ordinates_up)

    # 平滑化处理
    # v1
    # ordinates_low = ssp(ordinates_low[:, 0], ordinates_low[:, 1], False)
    # ordinates_up = ssp(ordinates_up[:, 0], ordinates_up[:, 1], True)

    # v2
    ordinates = ssp(ordinates[:, 0], ordinates[:, 1])

    # airfoil_show("ssp", ordinates, ordinates_low, ordinates_up, range_limit=False)

    # move y
    ordinates_low = move_y(ordinates_low, -ordinates_low[-1, 1])
    ordinates_up = move_y(ordinates_up, -ordinates_up[0, 1])

    # airfoil_show("move", ordinates, ordinates_low, ordinates_up)

    # rotate x
    # rotate low
    low_leading = ordinates_low[0]
    low_trailing = ordinates_low[-1]
    low_angle_vec = low_leading - low_trailing  ####  [0, 0]
    low_angle = - low_angle_vec[1] / low_angle_vec[0]  # negative
    ordinates_low = rotate1(ordinates_low, np.arctan(low_angle), low_trailing)

    # rotate up
    up_leading = ordinates_up[-1]
    up_trailing = ordinates_up[0]
    up_angle_vec = up_leading - up_trailing
    up_angle = - up_angle_vec[1] / up_angle_vec[0]  # positive
    ordinates_up = rotate1(ordinates_up, np.arctan(up_angle), up_trailing)

    # airfoil_show("rotate", ordinates, ordinates_low, ordinates_up)

    # normalize low;值得注意的是，下表面横坐标最小值不为0
    min = np.min(ordinates_low[:, 0])
    max = np.max(ordinates_low[:, 0])
    ordinates_low[:, 0] = (ordinates_low[:, 0] - min) / (max - min)

    # normalize up
    min = np.min(ordinates_up[:, 0])
    max = np.max(ordinates_up[:, 0])
    ordinates_up[:, 0] = (ordinates_up[:, 0] - min) / (max - min)
    # print("up:", ordinates_up)
    # print("low:", ordinates_low)

    airfoil_show("normalize", ordinates, ordinates_low, ordinates_up)

    # concatenate
    ordinates = np.concatenate([ordinates_up, ordinates_low[1::]], axis=0)
    # print(ordinates.shape)
    # print(ordinates)

    # plot
    # plt.title("concat")
    # plt.scatter(ordinates[:, 0], ordinates[:, 1])
    # plt.grid()  # 设置网格
    # plt.xlim((-0.05, 1.0))  # 设置横坐标范围
    # plt.ylim((-0.3, 0.3))  # 设置纵坐标范围
    # plt.xticks(np.arange(0, 1, 0.1))
    # plt.gca().set_aspect(1)
    # plt.show()

    # ordinates = pd.DataFrame(ordinates)
    # ordinates.columns = ['x', 'y']
    # ordinates.to_csv(os.path.join(save_dir, file), index=False)
    # print("save {}".format(file))


def file_to_ssp_v2(dir, file):
    '''
        按论文的方法，将弧长的开方代替横坐标
    :param dir:
    :param file:
    :return:
    '''
    file_path = os.path.join(dir, file)
    ordinates = pd.read_csv(file_path).to_numpy()

    # smooth.spline
    # 前缘顶点
    key = np.argmin(ordinates[:, 0])

    # 上表面坐标点
    ordinates_up = ordinates[0:key + 1]
    ordinates_up = ordinates_up[::-1]

    # 下表面坐标点
    ordinates_low = ordinates[key::]

    # 原翼型形状展示
    # airfoil_show("origin", ordinates, ordinates_low, ordinates_up)

    # 对坐标点进行一些预处理，将预处理后的结果再作平滑化
    arc_ordinates_up, arc_ordinates_low, arc_ordinates = x_to_arc(ordinates_up, ordinates_low, isShow=True)

    # 平滑化处理
    # v2：由于坐标点是一条曲线构成的，因此可以将所有坐标点一起进行平滑化，这正是我们的目的
    ssp_ordinates = sspv2(arc_ordinates[:, 0], arc_ordinates[:, 1])
    airfoil_show("ssp", ssp_ordinates, arc_ordinates_low, arc_ordinates_up, range_limit=True)

    # 把翼型进行恢复，只把纵坐标替换为平滑后的纵坐标，横坐标依然是最原始的
    re_ordinates = np.zeros_like(ordinates)
    re_ordinates[:, 0] = ordinates[:, 0]
    re_ordinates[:, 1] = ssp_ordinates[::-1, 1]

    # 平移整个翼型，使得前缘点落在(0,0)
    re_ordinates = move_y(re_ordinates, -re_ordinates[key, 1])
    airfoil_show2("move y", ordinates, re_ordinates)

    # 上表面坐标点
    re_ordinates_up = re_ordinates[0:key + 1]
    re_ordinates_up = re_ordinates_up[::-1]

    # 下表面坐标点
    re_ordinates_low = re_ordinates[key::]
    airfoil_show("recover1", np.asarray([[1, 0.05]]), re_ordinates_low, re_ordinates_up, range_limit=False)

    # 绕前缘点分别对上下翼型进行旋转，使得上下翼型后缘点落在(1,0)
    re_ordinates_up = rotate(re_ordinates_up, re_ordinates_up[-1])
    re_ordinates_low = rotate(re_ordinates_low, re_ordinates_low[-1])

    # red表示旋转后,blue表示旋转前
    airfoil_show("rotate", re_ordinates, re_ordinates_low, re_ordinates_up, range_limit=False)

    # 归一化
    # normalize low;值得注意的是，下表面横坐标最小值不为0
    min = np.min(re_ordinates_low[:, 0])
    max = np.max(re_ordinates_low[:, 0])
    re_ordinates_low[:, 0] = (re_ordinates_low[:, 0] - min) / (max - min)

    # normalize up
    min = np.min(re_ordinates_up[:, 0])
    max = np.max(re_ordinates_up[:, 0])
    re_ordinates_up[:, 0] = (re_ordinates_up[:, 0] - min) / (max - min)

    # concatenate
    re_ordinates = np.concatenate([re_ordinates_up[::-1], re_ordinates_low[1::]], axis=0)

    airfoil_show2("normalize", ordinates, re_ordinates)

    # re_ordinates = pd.DataFrame(re_ordinates)
    # re_ordinates.columns = ['x', 'y']
    # re_ordinates.to_csv(os.path.join(save_dir, file), index=False)
    # print("save {}".format(file))


def file_to_ssp_v3(dir, file):
    '''
        将下表面翼型关于y轴作对称(下表面翼型横坐标取反)
        与方法一区别不大，前缘也会变尖
    :param dir:
    :param file:
    :return:
    '''
    file_path = os.path.join(dir, file)
    ordinates = pd.read_csv(file_path).to_numpy()
    ordinates_tmp = ordinates.copy()

    # smooth.spline
    # 前缘顶点
    key = np.argmin(ordinates[:, 0])

    # 上表面坐标点
    ordinates_up = ordinates[0:key + 1]
    ordinates_up = ordinates_up[::-1]

    # 下表面坐标点
    ordinates_low = ordinates[key::]
    ordinates_low[:, 0] = -ordinates_low[:, 0]

    # 原翼型形状展示
    # airfoil_show("origin", ordinates, ordinates_low, ordinates_up)

    # 上下表面翼型合并
    cur_ordinates = np.concatenate([ordinates_low[::-1], ordinates_up[1::]], axis=0)

    # 平滑化处理
    # v2：由于坐标点是一条曲线构成的，因此可以将所有坐标点一起进行平滑化，这正是我们的目的
    ssp_ordinates = sspv2(cur_ordinates[:, 0], cur_ordinates[:, 1])
    # airfoil_show("ssp", ssp_ordinates, ordinates_low, ordinates_up, range_limit=True)

    # 把翼型进行恢复，只把纵坐标替换为平滑后的纵坐标，横坐标依然是最原始的
    re_ordinates = np.zeros_like(ordinates)
    re_ordinates[:, 0] = ordinates[:, 0]
    re_ordinates[:, 1] = ssp_ordinates[::-1, 1]

    # 平移整个翼型，使得前缘点落在(0,0)
    re_ordinates = move_y(re_ordinates, -re_ordinates[key, 1])
    # print("move dis: ", -re_ordinates[key, 1])
    # airfoil_show2("move y", ordinates, re_ordinates)

    # 上表面坐标点
    re_ordinates_up = re_ordinates[0:key + 1]
    re_ordinates_up = re_ordinates_up[::-1]

    # 下表面坐标点
    re_ordinates_low = re_ordinates[key::]
    # airfoil_show("recover1", np.asarray([[1, 0.05]]), re_ordinates_low, re_ordinates_up, range_limit=True)

    # 绕前缘点分别对上下翼型进行旋转，使得上下翼型后缘点落在(1,0)
    re_ordinates_up = rotate(re_ordinates_up, re_ordinates_up[-1])
    re_ordinates_low = rotate(re_ordinates_low, re_ordinates_low[-1])

    # red表示旋转后,blue表示旋转前
    airfoil_show("rotate", re_ordinates, re_ordinates_low, re_ordinates_up, range_limit=True)

    # 归一化
    # normalize low;值得注意的是，下表面横坐标最小值不为0
    re_ordinates_low[:, 0] = -re_ordinates_low[:, 0]
    min = np.min(re_ordinates_low[:, 0])
    max = np.max(re_ordinates_low[:, 0])
    re_ordinates_low[:, 0] = (re_ordinates_low[:, 0] - min) / (max - min)

    # normalize up
    min = np.min(re_ordinates_up[:, 0])
    max = np.max(re_ordinates_up[:, 0])
    re_ordinates_up[:, 0] = (re_ordinates_up[:, 0] - min) / (max - min)

    # concatenate
    re_ordinates = np.concatenate([re_ordinates_up[::-1], re_ordinates_low[-2::-1]], axis=0)

    airfoil_show2("normalize", ordinates_tmp, re_ordinates)
    # print(re_ordinates.shape)
    # print(re_ordinates)

    # re_ordinates = pd.DataFrame(re_ordinates)
    # re_ordinates.columns = ['x', 'y']
    # re_ordinates.to_csv(os.path.join(save_dir, file), index=False)
    # print("save {}".format(file))



# test
# file_to_ssp(dir, "ag04.csv")


# test x_to_arc
# ordinates_up = np.asarray([[0, 0], [0.25, 0.25], [0.5, 0.5], [1, 0]])
# ordinates_low = np.asarray([[0, 0], [0.25, -0.25], [0.5, -0.5], [1, 0]])
# ordinates_up, ordinates_low, ordinates = x_to_arc(ordinates_up, ordinates_low)



# read_dir = r"D:\data_cmp\airfoil_csv_op\NACA4"
# save_dir = r"D:\data_cmp\airfoil_csv_op\NACA4_ssp"
#
# for file in os.listdir(read_dir):
#     file_to_ssp_v3(read_dir, file)
#     break

